import gradio as gr
import torch
from PIL import Image
from torchvision import transforms
import timm
import torch.nn as nn
import numpy as np
import webbrowser

# 加载模型
def load_model(model_path, device, n_class):
    # 加载模型权重
    checkpoint = torch.load(model_path, map_location=device, weights_only=True)

    # 创建模型架构
    model = timm.create_model('resnest50d', pretrained=False)
    num_features = model.fc.in_features

    # 替换全连接层以匹配保存的权重
    model.fc = nn.Linear(num_features, n_class)  # n_class 必须与保存权重时的类别数量一致
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载权重
    model.to(device)
    model.eval()
    return model

# 图像预处理
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    image = image.convert('RGB')
    image_tensor = transform(image).unsqueeze(0)  # 添加批次维度
    return image_tensor

# 图像识别
def recognize_image(image, model, device, idx_to_labels):
    image_tensor = preprocess_image(image)
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        outputs = model(image_tensor)
        _, predicted = torch.max(outputs, 1)
        predicted_label = idx_to_labels[predicted.item()]
    return predicted_label

def open_handian(predicted_label):
    base_url = "https://www.zdic.net/"
    if predicted_label:
        search_url = f"{base_url}hans/{predicted_label}"
    else:
        search_url = base_url
    webbrowser.open_new_tab(search_url)
    return search_url  # 返回 URL 只是为了在 Gradio 界面上显示

if __name__ == '__main__':
    # 模型路径和类别映射路径
    model_path = 'best50_model.pth'  # 替换为您的模型路径
    idx_to_labels_path = 'idx_to_labels.npy'  # 替换为您的类别映射路径
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 确保 n_class 与保存权重时的类别数量一致
    n_class = 1781  # 根据保存的权重修改为正确的类别数量

    # 加载模型和标签映射
    model = load_model(model_path, device, n_class)
    idx_to_labels = np.load(idx_to_labels_path, allow_pickle=True).item()

    # 创建 Gradio 接口
    with gr.Blocks() as iface:
        gr.Markdown("## 甲骨文图像识别")
        image_input = gr.Image(type="pil", label="上传甲骨文图片")
        result_label = gr.Label(num_top_classes=1, label="识别结果")
        handian_url = gr.Textbox(label="汉典链接", visible=False)  # 用于存储汉典链接，不直接显示
        handian_button = gr.Button("查看汉典相关释义")

        def predict_and_get_url(image):
            if image is None:
                return "", "https://www.zdic.net/"
            predicted_label = recognize_image(image, model, device, idx_to_labels)
            handian_link = f"https://www.zdic.net/hans/{predicted_label}" if predicted_label else "https://www.zdic.net/"
            return predicted_label, handian_link

        output = image_input.change(
            fn=predict_and_get_url,
            inputs=image_input,
            outputs=[result_label, handian_url]
        )

        handian_button.click(
            fn=lambda url: webbrowser.open_new_tab(url) if url else None,
            inputs=handian_url,
            outputs=None
        )

        gr.Examples(
            examples=['train/安/G_0527_甲骨文0.png', 'train/北/G_0261_甲骨文0.png', "train/好/G_0450_甲骨文0.png", "train/白/G_1118_後2.24.9合3395.png"],
            inputs=image_input,
            label="示例图片"  # 自定义显示名称
        )

    # 启动 Gradio 界面
    iface.launch()